from torch.nn.functional import binary_cross_entropy_with_logits
from torch.nn import BCEWithLogitsLoss
from torch import tensor, Tensor


class BCEWithLogitsLoss(BCEWithLogitsLoss):
    def __init__(self, weight=None, size_average=None, reduce=None, reduction='mean', pos_weight=None):
        print("Instantiating BCEWithLogitsLoss")
        print(f"{pos_weight=}")
        if pos_weight is not None and not isinstance(pos_weight, Tensor):
            print(f"{pos_weight} is not a Tensor, is {type(pos_weight)}")
            if isinstance(pos_weight, int) or isinstance(pos_weight, float):
                pos_weight = [pos_weight]
            pos_weight = tensor(pos_weight) 
        super().__init__(
            weight=weight,
            size_average=size_average,
            reduce=reduce, 
            reduction=reduction, 
            pos_weight=pos_weight) 